"""
{'content': '在当前的技术背景下，随着大数据及人工智能技术的迅猛发展，基于大规模预训练语言模型的数字人驱动系统正逐渐成为研究热点。行业现状表明，尽管已有不少企业和研究机构涉足数字人领域，‘}

项目方案子目标写作
1.建立数据模型  子标题内容SubSchema(content)
2.建立系统和用户的提示词模板  system 输出结果output_template      human项目标题、全文摘要以及所提供的题目和子题目
3.创建规划者链函数
1）创建解析器JsonOutputParser(pydantic_object=SubSchema)
2）创建PromptTemplate,partial填充数据   prompts.partial(output_template=_parser.get_format_instructions())
3) 创建规划者链对象
"""

from pydantic import BaseModel,Field
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser

from sshcode.rag.model_utils import getLLM


class SubSchema(BaseModel):
    content:str=Field(description="写作的内容")

_system_prompt = """
您是一名优秀的方案段落写手，您擅长完成方案文档的某个部分的写作。
根据项目标题、全文摘要以及所提供的题目和子题目结构完成该子题目内容的编写。

要求：
您需要站在我们设计角度去写内容，字数越多越好。
内容中不要有序号。
不要有其它的表达样式符号。

输出：
{output_template}
"""

_human_prompt = """
项目标题：{project_name}
摘要：{abstract}
题目：{title}
子题目：{sub_title}
"""

def create_writer(llm):
    template = ChatPromptTemplate.from_messages([
        ("system", _system_prompt),
        ("human", _human_prompt)
    ])

    parser = JsonOutputParser(pydantic_object=SubSchema)
    template = template.partial(output_template=parser.get_format_instructions())

    return template | llm | parser

if __name__ == '__main__':
    llm = getLLM()
    agent = create_writer(llm)

    r = agent.invoke(
        {
            "project_name":"基于大模型的数字人驱动方案",
            "abstract":"本项目旨在设计和实施一种基于大规模预训练语言模型（以下简称“大模型”）的数字人智能驱动解决方案。该方案将涵盖关键技术的研究、平台建设和实际应用案例，旨在实现高精度的情感捕捉与自然语言交互功能。",
            "title":"项目背景与目标",
            "sub_title":"虚拟技术的发展现状",
        }
    )

    print(r)